// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/functions/prestosql/types/timestamp_with_time_zone_type.h>

using namespace kumo::pollux::test;

namespace kumo::pollux::functions {
namespace {

class MapSubsetTest : public test::FunctionBaseTest {
 public:
  template <typename T>
  void testFloatNaNs() {
    static const auto kNaN = std::numeric_limits<T>::quiet_NaN();
    static const auto kSNaN = std::numeric_limits<T>::signaling_NaN();

    // Case 1: Non-constant search keys.
    auto data = make_row_vector(
        {make_map_vector_from_json<T, int32_t>({
             "{1:10, NaN:20, 3:null, 4:40, 5:50, 6:60}",
             "{NaN:20}",
         }),
         make_array_vector<T>({{1, kNaN, 5}, {kSNaN, 3}})});

    auto expected = make_map_vector_from_json<T, int32_t>({
        "{1:10, NaN:20, 5:50}",
        "{NaN:20}",
    });
    auto result = evaluate("map_subset(c0, c1)", data);
    assertEqualVectors(expected, result);

    // Case 2: Constant search keys.
    data = make_row_vector(
        {make_map_vector_from_json<T, int32_t>({
             "{1:10, NaN:20, 3:null, 4:40, 5:50, 6:60}",
             "{NaN:20}",
         }),
         BaseVector::wrap_in_constant(2, 0, make_array_vector<T>({{1, kNaN, 5}}))});
    expected = make_map_vector_from_json<T, int32_t>({
        "{1:10, NaN:20, 5:50}",
        "{NaN:20}",
    });
    result = evaluate("map_subset(c0, c1)", data);
    assertEqualVectors(expected, result);

    // Case 3: Map with Complex type as key.
    // Map: { [{1, NaN,3}: 1, {4, 5}: 2], [{NaN, 3}: 3, {1, 2}: 4] }
    data = make_row_vector({
        make_map_vector(
            {0, 2},
            make_array_vector<T>({{1, kNaN, 3}, {4, 5}, {kSNaN, 3}, {1, 2}}),
            make_flat_vector<int32_t>({1, 2, 3, 4})),
        make_nested_array_vector_from_json<T>({
            "[[1, NaN, 3], [4, 5]]",
            "[[1, 2, 3], [NaN, 3]]",
        }),
    });
    expected = make_map_vector(
        {0, 2},
        make_array_vector_from_json<T>({
            "[1, NaN, 3]",
            "[4, 5]",
            "[NaN, 3]",
        }),
        make_flat_vector<int32_t>({1, 2, 3}));

    result = evaluate("map_subset(c0, c1)", data);
    assertEqualVectors(expected, result);
  }
};

TEST_F(MapSubsetTest, bigintKey) {
  auto data = make_row_vector({
      make_map_vector_from_json<int64_t, int32_t>({
          "{1:10, 2:20, 3:null, 4:40, 5:50, 6:60}",
          "{1:10, 2:20, 4:40, 5:50}",
          "{}",
          "{2:20, 4:40, 6:60}",
      }),
      make_array_vector_from_json<int64_t>({
          "[1, 3, 5]",
          "[1, 3, 5, 7]",
          "[3, 5]",
          "[1, 3]",
      }),
  });

  // Constant keys.
  auto result = evaluate("map_subset(c0, array_constructor(1, 3, 5))", data);

  auto expected = make_map_vector_from_json<int64_t, int32_t>({
      "{1:10, 3:null, 5:50}",
      "{1:10, 5:50}",
      "{}",
      "{}",
  });

  assertEqualVectors(expected, result);

  // Non-constant keys.
  result = evaluate("map_subset(c0, c1)", data);
  assertEqualVectors(expected, result);

  // Empty list of keys. Expect empty maps.
  result = evaluate("map_subset(c0, array_constructor()::bigint[])", data);

  expected = make_map_vector_from_json<int64_t, int32_t>({"{}", "{}", "{}", "{}"});

  assertEqualVectors(expected, result);
}

TEST_F(MapSubsetTest, varcharKey) {
  auto data = make_row_vector({
      make_map_vector_from_json<std::string, int32_t>({
          "{\"apple\": 1, \"banana\": 2, \"Cucurbitaceae\": null, \"date\": 4, \"eggplant\": 5, \"fig\": 6}",
          "{\"banana\": 2, \"orange\": 4}",
          "{\"banana\": 2, \"fig\": 4, \"date\": 5}",
      }),
      make_array_vector_from_json<std::string>({
          "[\"apple\", \"Cucurbitaceae\", \"fig\"]",
          "[\"apple\", \"Cucurbitaceae\", \"date\", \"eggplant\"]",
          "[\"fig\"]",
      }),
  });

  // Constant keys.
  auto result = evaluate(
      "map_subset(c0, array_constructor('apple', 'some very looooong name', 'fig', 'Cucurbitaceae'))",
      data);

  auto expected = make_map_vector_from_json<std::string, int32_t>({
      "{\"apple\": 1, \"Cucurbitaceae\": null, \"fig\": 6}",
      "{}",
      "{\"fig\": 4}",
  });

  assertEqualVectors(expected, result);

  // Non-constant keys.
  result = evaluate("map_subset(c0, c1)", data);
  assertEqualVectors(expected, result);

  // Empty list of keys. Expect empty maps.
  result = evaluate("map_subset(c0, array_constructor()::varchar[])", data);

  expected = make_map_vector_from_json<std::string, int32_t>({"{}", "{}", "{}"});

  assertEqualVectors(expected, result);
}

TEST_F(MapSubsetTest, arrayKey) {
  auto data = make_row_vector({
      make_map_vector(
          {0, 2},
          make_array_vector_from_json<int32_t>({
              "[1, 2, 3]",
              "[4, 5]",
              "[]",
              "[1, 2]",
          }),
          make_flat_vector<std::string>(
              {"apple", "orange", "Cucurbitaceae", "date"})),
      make_nested_array_vector_from_json<int32_t>({
          "[[1, 2, 3], [4, 5, 6]]",
          "[[1, 2, 3], []]",
      }),
  });

  auto result = evaluate("map_subset(c0, c1)", data);

  auto expected = make_map_vector(
      {0, 1},
      make_array_vector_from_json<int32_t>({
          "[1, 2, 3]",
          "[]",
      }),
      make_flat_vector<std::string>({"apple", "Cucurbitaceae"}));

  assertEqualVectors(expected, result);
}

TEST_F(MapSubsetTest, floatNaNs) {
  testFloatNaNs<float>();
  testFloatNaNs<double>();
}

TEST_F(MapSubsetTest, timestampWithTimeZone) {
  const auto keys = make_flat_vector<int64_t>(
      {pack(1, 1),
       pack(2, 2),
       pack(3, 3),
       pack(4, 4),
       pack(5, 5),
       pack(6, 6),
       pack(1, 7),
       pack(2, 8),
       pack(4, 9),
       pack(5, 10),
       pack(2, 11),
       pack(4, 12),
       pack(6, 13)},
      TIMESTAMP_WITH_TIME_ZONE());
  const auto values = make_nullable_flat_vector<int32_t>(
      {10, 20, std::nullopt, 40, 50, 60, 70, 80, 90, 100, 110, 120, 130});
  const auto maps = make_map_vector({0, 6, 10, 10}, keys, values);

  // Test map with TimestampWithTimeZone keys and constant second arg.
  const auto constLookup = BaseVector::wrap_in_constant(
      4,
      0,
      make_array_vector(
          {0},
          make_flat_vector<int64_t>(
              {pack(1, 1), pack(3, 2), pack(5, 3)},
              TIMESTAMP_WITH_TIME_ZONE())));
  const auto expectedKeys = make_flat_vector<int64_t>(
      {pack(1, 1), pack(3, 3), pack(5, 50), pack(1, 7), pack(5, 10)},
      TIMESTAMP_WITH_TIME_ZONE());
  const auto expectedValues =
      make_nullable_flat_vector<int32_t>({10, std::nullopt, 50, 70, 100});
  const auto expected =
      make_map_vector({0, 3, 5, 5}, expectedKeys, expectedValues);
  auto result =
      evaluate("map_subset(c0, c1)", make_row_vector({maps, constLookup}));

  assertEqualVectors(expected, result);

  // Test map with TimestampWithTimeZone keys and non-constant second arg.
  const auto lookupKeys = make_flat_vector<int64_t>(
      {pack(1, 1),
       pack(3, 3),
       pack(5, 5),
       pack(1, 10),
       pack(3, 12),
       pack(5, 13),
       pack(7, 14),
       pack(3, 15),
       pack(5, 16),
       pack(1, 17),
       pack(3, 18)},
      TIMESTAMP_WITH_TIME_ZONE());
  const auto lookup = make_array_vector({0, 3, 7, 9}, lookupKeys);

  result = evaluate("map_subset(c0, c1)", make_row_vector({maps, lookup}));
  assertEqualVectors(expected, result);

  // Test map with TimestampWithTimeZone wrapped in a complex type as keys.
  const auto mapsWithRowKeys =
      make_map_vector({0, 6, 10, 10}, make_row_vector({keys}), values);
  const auto lookupWithRowKeys =
      make_array_vector({0, 3, 7, 9}, make_row_vector({lookupKeys}));
  const auto expectedWithRowKeys = make_map_vector(
      {0, 3, 5, 5}, make_row_vector({expectedKeys}), expectedValues);

  result = evaluate(
      "map_subset(c0, c1)",
      make_row_vector({mapsWithRowKeys, lookupWithRowKeys}));
  assertEqualVectors(expectedWithRowKeys, result);
}

} // namespace
} // namespace kumo::pollux::functions
